// entities/do.go
package entities

import (
	"alan/exp3/src/system"
	"crypto/rand"
	"fmt"
	"log"
	"math/big"
)

// DataOwner 代表数据所有者
type DataOwner struct {
	skB          *big.Int             // 私钥 b
	params       *system.SystemParams // 系统公共参数
	originalData [][]*big.Int         // 存储原始数据块，以便生成审计请求时使用
}

// NewDataOwner 创建一个新的数据所有者，并完成KeyGen过程
func NewDataOwner(params *system.SystemParams) (*DataOwner, error) {
	b, err := rand.Int(rand.Reader, params.Q)
	if err != nil {
		return nil, err
	}
	if b.Cmp(big.NewInt(0)) == 0 {
		b = big.NewInt(1)
	}
	log.Printf("[DO] Generated private key b (bit len %d)", b.BitLen())

	params.Gb.Exp(params.G, b, params.P)
	log.Printf("[DO] Calculated public key Gb.")

	return &DataOwner{skB: b, params: params}, nil
}

// EncryptFile 核心算法：对文件进行分块和加密。
func (do *DataOwner) EncryptFile(fileData []byte, fileID string, blockSize int) (map[string][]*big.Int, error) {
	log.Printf("[DO] Starting encryption for file '%s'...", fileID)
	blocks := make([]*big.Int, 0)
	for i := 0; i < len(fileData); i += blockSize {
		end := min(i+blockSize, len(fileData))
		block := new(big.Int).SetBytes(fileData[i:end])
		block.Mod(block, do.params.Q)
		blocks = append(blocks, block)
	}
	do.originalData = append(do.originalData, blocks)
	log.Printf("[DO] File split into %d blocks, each value is mod Q.", len(blocks))

	encryptedDataMap := make(map[string][]*big.Int)
	q := do.params.Q
	for _, cspID := range do.params.A_IDs {
		encryptedBlocks := make([]*big.Int, len(blocks))
		for i, xi := range blocks {
			h1Input := fmt.Sprintf("%s||%s||%d", cspID, fileID, i)
			h2Input := fmt.Sprintf("%s||%s||%d", cspID, fileID, i)
			h1Val, h2Val := system.H1(h1Input, q), system.H2(h2Input, q)

			term1_val := new(big.Int).Add(xi, h1Val)
			term1 := new(big.Int).Mul(do.params.A_val, term1_val)

			b_mul_xi := new(big.Int).Mul(do.skB, xi)
			term2 := new(big.Int).Add(b_mul_xi, h2Val)

			yki := new(big.Int).Add(term1, term2)
			encryptedBlocks[i] = yki
		}
		encryptedDataMap[cspID] = encryptedBlocks
	}
	log.Printf("[DO] Encryption finished. y_ki are now large integers.")
	return encryptedDataMap, nil
}

// RequestAudit 核心算法：生成审计请求。
func (do *DataOwner) RequestAudit(fileIndex, challengedBlocks int, fileID string) (*system.AuditRequest, error) {
	log.Println("[DO] Generating audit request...")
	originalBlocks := do.originalData[fileIndex]
	totalBlocks := len(originalBlocks)

	indices := make([]int, challengedBlocks)
	var weights []*big.Int

	q := do.params.Q
	a := do.params.A_val

	retryCount := 0
	var y_check, h1_sum, h2_sum *big.Int

	for {
		weights = make([]*big.Int, challengedBlocks)
		for i := 0; i < challengedBlocks; i++ {
			idx, _ := rand.Int(rand.Reader, big.NewInt(int64(totalBlocks)))
			indices[i] = int(idx.Int64())
			e, _ := rand.Int(rand.Reader, q)
			weights[i] = e
		}

		y_check, h1_sum, h2_sum = new(big.Int), new(big.Int), new(big.Int)

		for _, cspID := range do.params.A_IDs {
			for l := 0; l < challengedBlocks; l++ {
				idx, weight := indices[l], weights[l]
				xi := originalBlocks[idx]

				h1Input := fmt.Sprintf("%s||%s||%d", cspID, fileID, idx)
				h2Input := fmt.Sprintf("%s||%s||%d", cspID, fileID, idx)
				h1Val, h2Val := system.H1(h1Input, q), system.H2(h2Input, q)

				b_mul_xi := new(big.Int).Mul(do.skB, xi)
				term_y := new(big.Int).Add(b_mul_xi, h2Val)
				y_check.Add(y_check, new(big.Int).Mul(weight, term_y))

				h1_sum.Add(h1_sum, new(big.Int).Mul(weight, h1Val))
				h2_sum.Add(h2_sum, new(big.Int).Mul(weight, h2Val))
			}
		}

		if y_check.Cmp(a) < 0 {
			log.Printf("[DO] Found valid weights after %d retries. Y < a is satisfied.", retryCount)
			break
		}
		retryCount++
		if retryCount > 100 {
			log.Fatalf("FATAL: Could not find valid weights after %d retries. Try increasing bit size of 'a'/'qPrime'.", retryCount)
		}
		log.Printf("[DO] Retry #%d: Generated weights are invalid (Y >= a).", retryCount)
	}

	mod_exp := do.params.P_minus_1

	sumH1_exp := new(big.Int).Mul(do.skB, h1_sum)
	sumH1_exp.Mod(sumH1_exp, mod_exp)

	sumH2_exp := new(big.Int).Mod(h2_sum, mod_exp)

	lambda1 := new(big.Int).Exp(do.params.G, sumH1_exp, do.params.P)
	lambda2 := new(big.Int).Exp(do.params.G, sumH2_exp, do.params.P)

	return &system.AuditRequest{
		FileID: fileID, Indices: indices, Weights: weights,
		Lambda1: lambda1, Lambda2: lambda2,
	}, nil
}
